package com.learn.security.utils.asymmetric_encryption;

import com.learn.security.utils.BaseUtil;

import javax.crypto.*;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

/**
 * DH加密工具类
 *
 * <pre>
 * DH
 * Default Keysize 1024
 * Keysize must be a multiple of 64, ranging from 512 to 1024 (inclusive).
 * </pre>
 */
public class DHUtil extends BaseUtil {

    public static final String ALGORITHM = "DH";

    /**
     * 默认密钥字节数
     */
    public static int KEY_SIZE = 1024;

    /**
     * DH加密下需要一种对称加密算法对数据加密，这里我们使用DES，也可以使用其他对称加密算法。
     */
    public static final String SECRET_ALGORITHM = "DES";
    public static final String PUBLIC_KEY = "DHPublicKey";
    public static final String PRIVATE_KEY = "DHPrivateKey";

    /**
     * 初始化甲方密钥
     *
     * @param keySize 密钥字节数
     */
    public static Map<String, Object> initKey(int keySize) {
        Map<String, Object> keyMap = new HashMap<>();

        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
            keyPairGenerator.initialize(keySize);

            KeyPair keyPair = keyPairGenerator.generateKeyPair();

            // 甲方公钥
            DHPublicKey publicKey = (DHPublicKey) keyPair.getPublic();

            // 甲方私钥
            DHPrivateKey privateKey = (DHPrivateKey) keyPair.getPrivate();

            keyMap.put(PUBLIC_KEY, publicKey);
            keyMap.put(PRIVATE_KEY, privateKey);
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }

        return keyMap;
    }

    /**
     * 初始化甲方密钥
     * 默认密钥字节数为1024
     */
    public static Map<String, Object> initKey() {
        return initKey(KEY_SIZE);
    }

    /**
     * 初始化乙方密钥
     *
     * @param key 甲方公钥
     */
    public static Map<String, Object> initKey(String key) {
        // 解析甲方公钥
        byte[] keyBytes = decryptBASE64(key);

        Map<String, Object> keyMap = new HashMap<>();

        try {
            if (null != keyBytes) {
                X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(keyBytes);
                KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
                PublicKey pubKey = keyFactory.generatePublic(x509KeySpec);

                // 由甲方公钥构建乙方密钥
                DHParameterSpec dhParamSpec = ((DHPublicKey) pubKey).getParams();

                KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
                keyPairGenerator.initialize(dhParamSpec);


                KeyPair keyPair = keyPairGenerator.generateKeyPair();

                // 乙方公钥
                DHPublicKey publicKey = (DHPublicKey) keyPair.getPublic();

                // 乙方私钥
                DHPrivateKey privateKey = (DHPrivateKey) keyPair.getPrivate();

                keyMap.put(PUBLIC_KEY, publicKey);
                keyMap.put(PRIVATE_KEY, privateKey);

            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return keyMap;
    }

    /**
     * 构建本地密钥
     *
     * @param publicKey  一方公钥
     * @param privateKey 另一方私钥
     * @return 本地秘钥
     */
    private static SecretKey getSecretKey(String publicKey, String privateKey) {
        // 初始化公钥
        byte[] pubKeyBytes = decryptBASE64(publicKey);
        // 初始化私钥
        byte[] priKeyBytes = decryptBASE64(privateKey);

        try {
            if (null != pubKeyBytes && null != priKeyBytes) {
                KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
                X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(pubKeyBytes);
                PublicKey pubKey = keyFactory.generatePublic(x509KeySpec);

                PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(priKeyBytes);
                Key priKey = keyFactory.generatePrivate(pkcs8KeySpec);

                KeyAgreement keyAgree = KeyAgreement.getInstance(ALGORITHM);
                keyAgree.init(priKey);
                keyAgree.doPhase(pubKey, true);

                // 生成本地密钥
                return keyAgree.generateSecret(SECRET_ALGORITHM);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    /**
     * 加密<br>
     *
     * @param data       待加密数据
     * @param publicKey  一方公钥--base64加密的字符
     * @param privateKey 另一方私钥--base64加密的字符
     * @return 加密后的byte[]
     */
    public static byte[] encrypt(byte[] data, String publicKey, String privateKey) {
        // 生成本地密钥
        SecretKey secretKey = getSecretKey(publicKey, privateKey);

        try {
            if (null != secretKey) {
                // 数据加密
                Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
                cipher.init(Cipher.ENCRYPT_MODE, secretKey);

                return cipher.doFinal(data);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 解密<br>
     *
     * @param data       待解密数据
     * @param publicKey  一方公钥--base64加密的字符
     * @param privateKey 另一方私钥--base64加密的字符
     * @return 解密后的byte[]
     */
    public static byte[] decrypt(byte[] data, String publicKey, String privateKey) {
        // 生成本地密钥
        SecretKey secretKey = getSecretKey(publicKey, privateKey);

        try {
            if (null != secretKey) {
                // 数据解密
                Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
                cipher.init(Cipher.DECRYPT_MODE, secretKey);

                return cipher.doFinal(data);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    /**
     * 取得私钥
     *
     * @return 私钥--base64加密的字符
     */
    public static String getPrivateKey(Map<String, Object> keyMap) {
        Key key = (Key) keyMap.get(PRIVATE_KEY);

        return encryptBASE64(key.getEncoded());
    }

    /**
     * 取得公钥
     *
     * @return 公钥--base64加密的字符
     */
    public static String getPublicKey(Map<String, Object> keyMap) {
        Key key = (Key) keyMap.get(PUBLIC_KEY);

        return encryptBASE64(key.getEncoded());
    }

}